# 导入所需的库函数
from ultralytics import YOLO
from PIL import Image
import os

# Load your model 导入模型
# model = YOLO('./model/yolov8n.pt')  # For example, 'yolov8n.pt'
model = YOLO('runs/detect/train2/weights/best.pt')  # For example, 'yolov8n.pt'
# 定义图片读取文件
path_input = "./image"
# 定义检测结果保存文件
path_output = "./output"

# 使用模型进行目标检测
for file_name in os.listdir(path_input):
    # 得到图片路径
    path_image = os.path.join(path_input, file_name)
    # Run prediction 预测
    results = model(path_image)
    # After running the input through the model, it returns an array of results for each input image.
    # As we provided only a single image, it returns an array with a single item that you can extract like this:
    result = results[0]
    # Now, iterate over detected objects
    for det in result.boxes:
        # det is now a single detection with attributes you can directly access
        xmin, ymin, xmax, ymax = det.xyxy[0]  # Coordinates
        conf = det.conf  # Confidence
        cls = det.cls  # Class ID
        class_name = result.names[cls[0].item()]   # class name
        print(f"Box coordinates: {xmin}, {ymin}, {xmax}, {ymax}, Confidence: {conf}, Class Name: {class_name}")
    # 显示图片并保存（这里如果使用matplotlib显示不出图片）
    image = Image.fromarray(result.plot()[:, :, ::-1])
    image.show()
    image.save(os.path.join(path_output, file_name))